# -*- coding:utf-8 -*-  
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
import datetime
import torch
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
from train import train_ED
from test import test
import config as c
from model.discriminator import Discriminator
from model.encoder_decoder import EncoderDecoder
from options import HiDDenConfiguration
from model.attack import Attack
import torch.optim as optim
import config as c
import time
from model_ED_N import ED_N
from tqdm import tqdm

def model_save(model, ep):
    save_path = c.model_dir
    print(f'--- save the model @ epoch {ep} ---')

    # 创建包含模型和优化器状态的 checkpoint
    checkpoint = {
        'epoch': ep,
        'encoder_decoder_state_dict': model.encoder_decoder.state_dict(),  # 模型参数
        'dis_state_dict' : model.discriminator.state_dict(),
        'noiser_state_dict': model.blockII.state_dict(),
        'ed_opt_state_dict': model.optimizer_ed.state_dict(),  # 优化器状态
        'dis_opt_state_dict':model.optimizer_dis.state_dict(),
        'noiser_opt_state_dict':model.optimizer_n.state_dict()
    }

    # 保存 checkpoint
    torch.save(checkpoint, f'{save_path}/{ep:05d}.pth')


def train(model, train_loader, test_loader):
    
    for ep in range(model.epoch, c.epochs):
        model.res_init()  # 初始化模型
        start_time = time.time()  # Start time of the epoch
        
        train_single_epoch(ep+1, model, train_loader)

        # 保存模型
        # model_save(model, ep+1)

        # 测试模型
        model.test_EDN(ep+1, len(train_loader), test_loader)

        end_time = time.time()  # End time of the epoch

        epoch_duration = end_time - start_time
        print(f"Epoch {ep+1} took {epoch_duration:.2f} seconds")

def train_single_epoch(ep, model, train_loader):

    for i, (images, _) in enumerate(tqdm(train_loader, desc=f"Epoch {ep}/{c.epochs}", unit="batch")):
        images = images.cuda()
        model.update_encoder_decoder_ASL(images)
        model.update_encoder_decoder_SM(images)
        model.update_noiser(images)       

def main():
    torch.autograd.set_detect_anomaly(True)
    file_dir = "runs-works"
    if not os.path.exists(file_dir):
        os.makedirs(file_dir)
    # 获取当前日期和时间
    current_time = datetime.datetime.now()
    # 格式化日期和时间
    time_str = current_time.strftime("%Y-%m-%d_%H-%M-%S")

    # name = "Brightness_sampled_1000_en0.7_de1.0" + "_" + time_str
    # name = "Contrast_sampled_1000_en0.7_de1.0" + "_" + time_str
    # name = "Saturation_sampled_1000_en0.7_de1.0" + "_" + time_str
    # name = "Hue_sampled_1000_en0.7_de1.0" + "_" + time_str
    # name = "Jpeg_sampled_1000_en0.7_de1.0" + "_" + time_str

    # name = "Gaussain_sampled_1000_en0.7_de1.0_Net100" + "_" + time_str
    # name = "Salt_sampled_1000_en0.7_de1.0_Net100" + "_" + time_str

    # name = "Gaussain_sampled_1000_en0.7_de1.0_tanh" + "_" + time_str
    # name = "Gaussain_sampled_1000_en0.7_de1.0_sigmoid" + "_" + time_str
    name = "Salt_sampled_1000_en0.7_de1.0_tanh" + "_" + time_str
    # name = "Salt_sampled_1000_en0.7_de1.0_sigmoid" + "_" + time_str
    c.file_name = name

    train_dir = os.path.join(file_dir, name)
    os.makedirs(train_dir, exist_ok=True)


    model_dir = "model"
    c.model_dir = os.path.join(train_dir, model_dir)
    #runs-works/name/model
    os.makedirs(c.model_dir, exist_ok=True)
    
    log_dir = "log"
    log_dir = os.path.join(train_dir, log_dir)
    
    os.makedirs(log_dir, exist_ok=True)

    c.log_file = os.path.join(log_dir, "log.csv")

    save_image_dir = "image"
    c.save_image_dir = os.path.join(train_dir, save_image_dir)
    os.makedirs(c.save_image_dir, exist_ok=True)

    data_dir = "/home/zzc/simulator_noisy/sampled_data"
    # sample_data_dir = data_dir + "/sample_100"
    # sample_data_dir = data_dir + "/sample_500"
    sample_data_dir = data_dir + "/sample_1000"
    # sample_data_dir = data_dir + "/sample_5000"

    train_data_dir = sample_data_dir + "/train"  # 训练数据
    test_data_dir = sample_data_dir + "/val"     # 测试数据


    # 定义图像预处理的转换操作
    data_transforms = {
        'train': transforms.Compose([
            transforms.RandomCrop((c.image_size, c.image_size), pad_if_needed=True),
            transforms.ToTensor()
        ]),
        'test': transforms.Compose([
            transforms.CenterCrop((c.image_size, c.image_size)),
            transforms.ToTensor()
        ])
    }

    # 创建ImageFolder数据集实例
    train_dataset = ImageFolder(train_data_dir, transform=data_transforms['train'])
    test_dataset = ImageFolder(test_data_dir, transform=data_transforms['test'])

    # 创建数据加载器
    train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=4, drop_last=True)
    test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False, num_workers=4, drop_last=True)
    
    model = ED_N().cuda()
    train(model, train_loader, test_loader)

if __name__ == '__main__':
    main()